import numpy as np


def forward_kinematics(DH_params):
    # DH_params: DH参数矩阵，每行代表一个连杆的DH参数
    # T: 末端执行器的变换矩阵

    # 初始化变换矩阵为单位矩阵
    T = np.eye(4)

    # 遍历所有的DH参数
    for i in range(DH_params.shape[0]):
        # 将角度转为弧度
        alpha = np.deg2rad(DH_params[i, 1])
        a = DH_params[i, 2]
        d = DH_params[i, 3]
        theta = np.deg2rad(DH_params[i, 4])

        # 计算当前连杆的变换矩阵
        A_i = np.array([
            [np.cos(theta), -np.sin(theta) * np.cos(alpha), np.sin(theta) * np.sin(alpha), a * np.cos(theta)],
            [np.sin(theta), np.cos(theta) * np.cos(alpha), -np.cos(theta) * np.sin(alpha), a * np.sin(theta)],
            [0, np.sin(alpha), np.cos(alpha), d],
            [0, 0, 0, 1]
        ])

        # 累乘到总变换矩阵T
        T = T @ A_i  # 使用@符号进行矩阵乘法

    return T

if __name__ == '__main__':
# 示例DH参数矩阵
    DH_params = np.array([
        [1, -45, 0, 560, 10],
        [2, 45, 0, 450, 0],
        [3, 45, 0, 410, 10],
        [4, -45, 0, 380, -20],
        [5, 45, 0, 360, 0],
        [6, 0, 0, 210, 40]
    ])

    # 调用函数计算变换矩阵
    T = forward_kinematics(DH_params)
    print(T)